#!/usr/bin/perl -w
use strict;
use warnings;

if (@ARGV != 2) {
  print "USAGE: $0 codonfreq.txt output.fasta\n";
  exit(1);
}

my ($codonTableFile, $fastaFile) = @ARGV; 

my %codonTable = getCodonTableFrom ($codonTableFile);
my @seqs = getFastaFrom ($fastaFile);

if (@seqs == 0) {
  print "Need at least one sequence in $fastaFile.\n";
  exit(1);
}

foreach my $seq (@seqs) {
  # get the coding sequence
  $seq =~ s/^.*?([ACGT]{3}.*?[ACGT]{3}).*$/$1/g;
  $seq = uc $seq;
}

my $stddev = codonFreqStdDev(\@seqs, \%codonTable);

printf "Standard Deviation: %.5f\n", $stddev;
printf "This means that, on average, every codon is being over- or under-represented by about %.1f percent.\n", $stddev*100/deviationFraction(1, %codonTable);


sub deviationFraction {
  my ($fraction, %codonTable) = @_;
  my $sum=0;
  foreach my $codon (keys %codonTable) {
    $sum += $codonTable{$codon}**2;
  }
  my $variance=$sum/64;
  return $fraction*sqrt($variance);
}

sub codonFreqStdDev {
  my ($seqs, $codonTable) = @_;
  my @seqs = @{$seqs};
  my %codonTable = %{$codonTable};
  my $sum=0;
  my $numCodons=0;
  my %codonCounts;

  # initialize counts to 0
  foreach my $codon (keys %codonTable) {
    $codonCounts{$codon} = 0;
  }

  foreach my $seq (@seqs) {
    # split the sequence into triplets of characters (codons)
    my @codons = ( $seq =~ m/.{3}/g );

    $numCodons += @codons;

    # gather counts
    foreach my $codon (@codons) {
      $codonCounts{$codon} += 1;
    }
  }

  # calculate deviations, add to $sum
  foreach my $codon (keys %codonCounts) {
      my $observed = $codonCounts{$codon}/$numCodons;
      my $expected = $codonTable{$codon};
      $sum += ($observed-$expected)**2;
  }

  my $variance = $sum/64; # 64 codons
  return sqrt($variance);
}

# parses a codon frequency table from the Codon Usage Database
sub getCodonTableFrom {
  my ($codonTableFile) = @_;
  my %codonTable;
  open (INPUT, $codonTableFile) or die "Can't open $codonTableFile: $!";

  while( my $line = <INPUT> ) {
    if ( $line =~ /^\s*$/ ) {
      next;
    }
    $line =~ s/\(.*?\)//g; # get rid of parens and the stuff inside of them
    $line = uc $line;
    $line =~ tr/U/T/; # convert all U's to T's (we're expecting DNA)
    my @stats = split (/\s+/, $line); # split on whitespace
    for (my $i=0; $i<4; $i++) {
      # divide by 1000 so that the sum of the frequencies is 1
      $codonTable{$stats[2*$i]} = $stats[2*$i+1]/1000; 
    }
  }
  return %codonTable;
}


# given a FASTA file, return an array of all sequences from that file (not 
# including sequence labels)
sub getFastaFrom {
  my ($fastafile) = @_;
  my $qty = 0;
  my @seqs;
  open (INPUT, $fastafile) or die "Can't open $fastafile: $!";

  while( my $line = <INPUT> ) {
    if( $line =~ /^\s*$/ ) {
      next;
    } elsif( $line =~ /^>/ ) {
      # increment the quantity of sequences if a label is identified
      $qty++;
    } else {
      # concatenate the line to the current sequence if it is not a label
      $line =~ s/[^ACGT]//gi;
      $seqs[$qty-1] .= $line;
    }
  }
  return @seqs;
}

